library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)

if (!file.exists("output")) {
  dir.create("output")
}

# load data
source("HNC_metadata_tidy.R")

sbsCOSMIC <- read.csv("Input_signature_attributions/Input_SBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>%
  tibble::rownames_to_column("donor_id") %>%
  mutate(
    APOBEC = SBS2 + SBS13,
    others = SBS7a + SBS7b + SBS7c + SBS15 + SBS17a + SBS17b + SBS33 + SBS93 + SBS_L
  ) %>%
  select(donor_id, APOBEC, SBS1, SBS4, SBS5, SBS12, SBS16, SBS18, SBS39, SBS92, SBS_I, others)

dbsCOSMIC <- read.csv("Input_signature_attributions/Input_DBS_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")
idCOSMIC <- read.csv("Input_signature_attributions/Input_ID_COSMIC_attributions.csv", stringsAsFactors = T, row.names = 1) %>% tibble::rownames_to_column("donor_id")

additional_df <- list(sbsCOSMIC, dbsCOSMIC, idCOSMIC)

# merge with cosmic sigantures and generate new tobacco + hpv variable
for (df in additional_df) {
  data <- data %>% left_join(df, by = "donor_id")
}

data <- data %>%
  mutate(
    tob_hpv = case_when(
      tobacco != "Never" & hpv_pos == "Positive" ~ "tobacco + HPV",
      tobacco != "Never" & hpv_pos == "Negative" ~ "tobacco",
      tobacco == "Never" & hpv_pos == "Positive" ~ "HPV",
      tobacco == "Never" & hpv_pos == "Negative" ~ "none"
    ),
    tob_hpv = factor(tob_hpv, levels = c("HPV", "tobacco + HPV", "none", "tobacco"))
  )

# Calculate mean signature burden
factors <- c("hpv_pos", "tob_hpv")
sigtype <- sbsCOSMIC
relative <- TRUE

p <- list()
# Subset of orpoharynx samples
data.s <- data %>% filter(donor_id %in% sigtype$donor_id, subsite == "OPC")
for (f in factors) {
  cohort_average <- data.s %>%
    group_by(data.s[, f]) %>%
    summarise_at(vars(names(sigtype)[-1]), mean) %>%
    na.omit()
  cohort_average <- data.frame(cohort_average, row.names = 1)
  cohort_average_plot <- as.data.frame(t(cohort_average))
  cohort_average_plot <- tibble::rownames_to_column(cohort_average_plot, "Signature")
  cohort_average_plot_melted <- gather(cohort_average_plot, variable, value, -Signature)
  cohort_average_plot_melted$Signature <- factor(cohort_average_plot_melted$Signature,
                                                 levels = names(sigtype)[-1]
  )
  p <- ggplot(cohort_average_plot_melted, aes(fill = Signature, x = variable, y = value)) +
    geom_bar(position = ifelse(relative == T, "fill", "stack"), stat = "identity", colour = "black", width = 1) +
    scale_x_discrete(limits = levels(data.s[, f])) +
    scale_fill_viridis(discrete = T, direction = -1) +
    theme_bw() +
    ylab(ifelse(relative == T, "Average Relative Proportion", "Average burden")) +
    xlab(ifelse(f == "tob_hpv", "Tobacco and HPV status", f)) +
    ggtitle(paste("Signature distribution")) +
    theme(
      plot.title = element_text(face = "bold", size = 12),
      axis.text.x = element_text(size = 12, angle = 45, hjust = 1),
      axis.text.y = element_text(size = 12),
      axis.title.y = element_text(size = 12),
      legend.text = element_text(size = 12),
      legend.title = element_blank(),
      legend.position = "bottom"
    )
  print(p)
  ggsave(
    plot = p, paste0("output/ExtendedDataFigure_7c_", f, "_", Sys.Date(), ".pdf"), device = "pdf",
    width = 5, height = 6
  )
}

